# client_demo.py
import asyncio
from mcp import ClientSession
from contextlib import AsyncExitStack
from mcp.client.sse import sse_client

class MCPClient:
    def __init__(self):
        self.session = None
        self.exit_stack = AsyncExitStack()

    async def connect_to_server(self, server_url: str):
        async with sse_client("http://localhost:8000/sse") as streams:
            async with ClientSession(streams[0], streams[1]) as session:
                try:
                    self.session = await session.initialize()
                except:
                    print('')

        """连接 MCP 服务端（Windows 下使用 SSE 协议）"""
        # with sse_client("http://localhost:8000/sse") as streams:
        #     with ClientSession(streams[0], streams[1]) as session:
        #         self.session = session.initialize()
        #     # self.session = await self.exit_stack.enter_async_context(
        #     #     ClientSession(server_url)
        #     # )
        #     print(f"✅ 已连接到 MCP 服务端：{server_url}")

    async def list_tools(self):
        """获取服务端注册的工具列表"""
        tools = await self.session.get_tools()
        print("\n可用工具列表：")
        for tool in tools:
            print(f"- 工具名：{tool.name}，描述：{tool.description}")

    async def call_tool(self, tool_name: str, **kwargs):
        """调用指定工具"""
        result = await self.session.call_tool(tool_name, **kwargs)
        print(f"\n工具 {tool_name} 调用结果：{result}")

    async def run(self):
        # 连接服务端（示例连接本地 8000 端口的文件服务）
        await self.connect_to_server("http://localhost:8000/sse")

        # 列出工具并调用
        await self.list_tools()
        await self.call_tool("get_desktop_files")  # 调用文件工具


if __name__ == "__main__":
    # Windows 下需使用 asyncio 的 ProactorEventLoop
    # asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
    client = MCPClient()
    asyncio.run(client.run())